package com.ls.lsojcodesandbox;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.command.ExecStartResultCallback;
import com.ls.lsojcodesandbox.model.ExecuteMessage;
import com.ls.lsojcodesandbox.utils.JavaCodeSandboxTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

@Component
public class JavaDockerSandbox extends JavaCodeSandboxTemplate {
    private static final String IMAGE_NAME = "openjdk:8-alpine";
    private static final long TIME_OUT = 5000L;
    private static final int CONTAINER_POOL_SIZE = 2; // 定义容器池大小

    private DockerClient dockerClient;
    private BlockingQueue<String> containerPool;

    /**
     * 服务启动时执行，用于初始化容器池
     */
    @PostConstruct
    private void init() {
        dockerClient = DockerClientBuilder.getInstance().build();
        containerPool = new LinkedBlockingQueue<>(CONTAINER_POOL_SIZE);

        // 1. 检查并拉取镜像
        pullImageIfNeeded();

        // 2. 创建并启动容器，放入池中
        for (int i = 0; i < CONTAINER_POOL_SIZE; i++) {
            CreateContainerResponse container = createContainer();
            dockerClient.startContainerCmd(container.getId()).exec();
            containerPool.add(container.getId());
            System.out.println("Container started and added to pool: " + container.getId());
        }
    }

    /**
     * 服务关闭时执行，用于清理容器
     */
    @PreDestroy
    private void cleanup() {
        System.out.println("Cleaning up containers...");
        for (String containerId : containerPool) {
            try {
                dockerClient.stopContainerCmd(containerId).exec();
                dockerClient.removeContainerCmd(containerId).exec();
            } catch (Exception e) {
                System.err.println("Failed to stop/remove container " + containerId + ": " + e.getMessage());
            }
        }
    }

    private void pullImageIfNeeded() {
        List<Image> images = dockerClient.listImagesCmd().withImageNameFilter(IMAGE_NAME).exec();
        if (images.isEmpty()) {
            System.out.println("Image not found locally, pulling: " + IMAGE_NAME);
            PullImageResultCallback callback = new PullImageResultCallback() {
                @Override
                public void onNext(PullResponseItem item) {
                    super.onNext(item);
                }
            };
            try {
                dockerClient.pullImageCmd(IMAGE_NAME).exec(callback).awaitCompletion();
                System.out.println("Image pulled successfully.");
            } catch (InterruptedException e) {
                System.err.println("Failed to pull image.");
                throw new RuntimeException(e);
            }
        } else {
            System.out.println("Image already exists: " + IMAGE_NAME);
        }
    }

    private CreateContainerResponse createContainer() {
        HostConfig hostConfig = new HostConfig();
        hostConfig.withCpuCount(1L)
                .withMemory(100 * 1000L * 1000L);
        hostConfig.setBinds(new Bind("/home/ls/JavaProject/lsojcodesandbox/tmpCode", new Volume("/app")));
        return dockerClient.createContainerCmd(IMAGE_NAME)
                .withHostConfig(hostConfig)
                .withNetworkDisabled(true) // 增强安全性，禁止网络访问
                .withAttachStdout(true)
                .withAttachStderr(true)
                .withTty(true)
                .withReadonlyRootfs(true) // 增强安全性，设置容器根文件系统为只读
                .exec();
    }


    @Override
    public List<ExecuteMessage> runFile(File userCodeFile, List<String> inputList, String lastFileName) {
        String containerId = null;
        try {
            // 1. 从池中获取一个容器
            containerId = containerPool.take();
            System.out.println("Took container from pool: " + containerId);

            // 2. 将用户代码拷贝到容器中
//            String userCodeParentPath = userCodeFile.getParent();
//            dockerClient.copyArchiveToContainerCmd(containerId)
//                    .withHostResource(userCodeParentPath)
//                    .withRemotePath("/app")
//                    .exec();
//            System.out.println("Copied files to container: " + containerId);

            // 3. 循环执行测试用例
            List<ExecuteMessage> executeMessages = new ArrayList<>();
            for (String input : inputList) {
                ExecuteMessage executeMessage = executeWithInput(containerId, input, lastFileName);
                executeMessages.add(executeMessage);
            }

            return executeMessages;

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("Failed to get container from pool", e);
        } finally {
            // 4. 将容器归还到池中
            if (containerId != null) {
                // 可选：在这里执行清理操作，例如 `docker exec rm -rf /app/*`
                containerPool.add(containerId);
                System.out.println("Returned container to pool: " + containerId);
            }
        }
    }

    private ExecuteMessage executeWithInput(String containerId, String input, String lastFileName) {
        // 1. (推荐) 在容器内创建子目录，-p 能确保目录存在且不报错
        String mkdirCommand = String.format("mkdir -p /app/%s", lastFileName);
        //execCmd(containerId, new String[]{"sh", "-c", mkdirCommand});


        // 2. 构造新的输入文件路径并执行写入命令
        String inputFilePathInContainer = String.format("/app/%s/input.txt", lastFileName);
        String writeInputCommand = "echo '" + input.replace("'", "'\\''") + "' > " + inputFilePathInContainer;
        execCmd(containerId, new String[]{"sh", "-c", writeInputCommand});


        // 3. 构造新的 Java 执行命令
        // 新的 Classpath 目录
        String subDirPathInContainer = String.format("/app/%s", lastFileName);
        // 完整的 java 命令，注意 -cp 和重定向路径都已更新
        String javaCommand = String.format("java -cp %s Main < %s", subDirPathInContainer, inputFilePathInContainer);
        // 最终在容器中执行的命令数组
        String[] cmdArray = {"sh", "-c", javaCommand};        ExecCreateCmdResponse execCreateCmdResponse = dockerClient.execCreateCmd(containerId)
                .withCmd(cmdArray)
                .withAttachStderr(true)
                .withAttachStdout(true)
                .exec();

        ExecuteMessage executeMessage = new ExecuteMessage();
        final long[] maxMemory = {0L};
        final boolean[] timeout = {true}; // 默认未超时

        // 监控内存
        StatsCmd statsCmd = dockerClient.statsCmd(containerId);
        ResultCallback<Statistics> statsCallback = statsCmd.exec(new ResultCallback<Statistics>() {
            @Override
            public void onStart(Closeable closeable) {}
            @Override
            public void onNext(Statistics stats) {
                maxMemory[0] = Math.max(maxMemory[0], stats.getMemoryStats().getUsage());
            }
            @Override
            public void onError(Throwable throwable) {}
            @Override
            public void onComplete() {}
            @Override
            public void close() {}
        });

        // 获取输出
        final StringBuilder messageBuilder = new StringBuilder();
        final StringBuilder errorBuilder = new StringBuilder();
        ExecStartResultCallback execStartResultCallback = new ExecStartResultCallback() {
            @Override
            public void onNext(Frame frame) {
                StreamType streamType = frame.getStreamType();
                if (StreamType.STDERR.equals(streamType)) {
                    errorBuilder.append(new String(frame.getPayload()));
                } else {
                    messageBuilder.append(new String(frame.getPayload()));
                }
            }
            @Override
            public void onComplete() {
                timeout[0] = false; // 正常完成，未超时
                super.onComplete();
            }
        };

        try {
            StopWatch stopWatch = new StopWatch();
            stopWatch.start();

            dockerClient.execStartCmd(execCreateCmdResponse.getId())
                    .exec(execStartResultCallback)
                    .awaitCompletion(TIME_OUT, TimeUnit.MILLISECONDS);

            stopWatch.stop();
            executeMessage.setTime(stopWatch.getLastTaskTimeMillis());

        } catch (InterruptedException e) {
            throw new RuntimeException("Execution interrupted", e);
        } finally {
            //停止监控
            try {
                statsCallback.close();
            } catch (IOException e) {
                // ignore
            }
        }

        executeMessage.setMemory(maxMemory[0]);
        executeMessage.setMessage(messageBuilder.toString());
        executeMessage.setErrorMessage(errorBuilder.toString());

        if (timeout[0]) {
            executeMessage.setErrorMessage("Execution timed out.");
        }

        return executeMessage;
    }

    private void execCmd(String containerId, String[] cmd) {
        ExecCreateCmdResponse execCreateCmdResponse = dockerClient.execCreateCmd(containerId)
                .withCmd(cmd)
                .exec();
        try {
            dockerClient.execStartCmd(execCreateCmdResponse.getId())
                    .exec(new ExecStartResultCallback())
                    .awaitCompletion();
        } catch (InterruptedException e) {
            throw new RuntimeException("Exec command failed", e);
        }
    }
}